

#include <errno.h>
#include <liburing.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define EVENT_ACCEPT 0
#define EVENT_READ 1
#define EVENT_WRITE 2
typedef struct uring_io_s uring_io_t;
typedef int (*data_handle)(uring_io_t *io, void *data, int datelen);

struct conn_info {
  int fd;
  int event;
};

typedef struct uring_loop_s {
  struct io_uring ring;
  struct conn_info accept_info;
} uring_loop_t;

struct uring_io_s {
  struct conn_info c_info;
  unsigned char *wbuff;
  unsigned char *rbuff;
  int wbuff_len;
  int rbuff_len;
  data_handle readcb;
  data_handle writecb;
  uring_loop_t *loop;
};

static int alloc_io_rwbuff(uring_io_t *io,int wbuff_len,int rbuff_len){
    io->wbuff_len = wbuff_len;
    io->rbuff_len = rbuff_len;
    io->wbuff = (unsigned char *)calloc(1, io->wbuff_len);
    if(!io->wbuff){
        return  -ENOMEM;
    }
    io->rbuff = (unsigned char *)calloc(1,io->rbuff_len);
    if(!io->rbuff){
        free(io->wbuff);
        return -ENOMEM;
    }
    return  0;
}

static int free_io_rwfree(uring_io_t *io){
    if(io->wbuff){
        free(io->wbuff);
    }
    if(io->rbuff){
        free(io->rbuff);
    }
    return 0;
}

int init_server(unsigned short port) {

  int sockfd = socket(AF_INET, SOCK_STREAM, 0);
  struct sockaddr_in serveraddr;
  memset(&serveraddr, 0, sizeof(struct sockaddr_in));
  serveraddr.sin_family = AF_INET;
  serveraddr.sin_addr.s_addr = htonl(INADDR_ANY);
  serveraddr.sin_port = htons(port);

  int opt = 1;
  setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

  if (0 <
      bind(sockfd, (struct sockaddr *)&serveraddr, sizeof(struct sockaddr))) {
    perror("bind");
    return -1;
  }

  listen(sockfd, 10);

  return sockfd;
}

#define ENTRIES_LENGTH 1024
#define BUFFER_LENGTH 1024

int set_event_recv(uring_io_t *io, int sockfd, void *buf, size_t len,
                   int flags) {

  struct io_uring_sqe *sqe = io_uring_get_sqe(&io->loop->ring);

  io_uring_prep_recv(sqe, sockfd, buf, len, flags);
  io_uring_sqe_set_data(sqe, io);
  return 0;
}

int set_event_send(uring_io_t *io, int sockfd, void *buf, size_t len,
                   int flags) {

  struct io_uring_sqe *sqe = io_uring_get_sqe(&io->loop->ring);

  io_uring_prep_send(sqe, sockfd, buf, len, flags);
  io_uring_sqe_set_data(sqe, io);
  return 0;
}

int set_event_accept(uring_loop_t *loop, int sockfd, struct sockaddr *addr,
                     socklen_t *addrlen, int flags) {
  struct io_uring_sqe *sqe = io_uring_get_sqe(&loop->ring);

  io_uring_prep_accept(sqe, sockfd, (struct sockaddr *)addr, addrlen, flags);
  io_uring_sqe_set_data(sqe, &loop->accept_info);
  return 0;
}

int main(int argc, char *argv[]) {

  unsigned short port = 12345;
  int sockfd = init_server(port);

  struct io_uring_params params;
  memset(&params, 0, sizeof(params));

  uring_loop_t loop;
  // io_uring_queue_init_params(ENTRIES_LENGTH, &ring, &params);
  int ret = io_uring_queue_init(ENTRIES_LENGTH, &loop.ring, 0);
  loop.accept_info.event = EVENT_ACCEPT;
  loop.accept_info.fd = sockfd;

#if 0
	struct sockaddr_in clientaddr;	
	socklen_t len = sizeof(clientaddr);
	accept(sockfd, (struct sockaddr*)&clientaddr, &len);
#else
#if 0
	struct sockaddr_in clientaddr;	
	socklen_t len = sizeof(clientaddr);
	set_event_accept(&loop, sockfd, (struct sockaddr*)&clientaddr, &len, 0);
#else
  set_event_accept(&loop, sockfd, NULL, 0, 0);
#endif

#endif

  char buffer[BUFFER_LENGTH] = {0};
    int r_counts = 0;
  while (1) {

    io_uring_submit(&loop.ring);

    struct io_uring_cqe *cqe;
    io_uring_wait_cqe(&loop.ring, &cqe);

    printf("rw_conts:%d\n",r_counts);
    struct io_uring_cqe *cqes[128];
    int nready = io_uring_peek_batch_cqe(&loop.ring, cqes, 128); // epoll_wait

    int i = 0;
    struct io_uring_cqe *entries = NULL;
    struct conn_info *result = NULL;
    uring_io_t *io = NULL;
    for (i = 0; i < nready; i++) {

      entries = cqes[i];
      result = (struct conn_info *)io_uring_cqe_get_data(entries);

      if (result->event == EVENT_ACCEPT) {

        set_event_accept(&loop, sockfd, NULL, 0, 0);
        // printf("set_event_accept\n"); //
        io = (uring_io_t *)calloc(1, sizeof(uring_io_t));
        if(!io){
            printf("alloc io error\n");
            continue;
        }
        // int connfd = entries->res;
        if(alloc_io_rwbuff(io, BUFSIZ, BUFSIZ)){
            free(io);
            io = NULL;
            continue;
        }
        io->loop = &loop;
        io->c_info.event = EVENT_READ;
        io->c_info.fd = entries->res;
        set_event_recv(io, entries->res, io->rbuff, io->rbuff_len, 0);

      } else if (result->event == EVENT_READ) { //

        ret = entries->res;
        io = io_uring_cqe_get_data(entries);
        // printf("set_event_recv ret: %d, fd:%d ,data:%s\n", ret, io->c_info.fd,io->rbuff); //
        if (ret <= 0) {
          printf("rclose fd:%d,ret:%d\n",io->c_info.fd,ret);
            close(result->fd);
          free_io_rwfree(io);
          free(io);
        } else if (ret > 0) {
          io->c_info.event = EVENT_WRITE;
            ++r_counts;
          //   printf("read\n");
          set_event_send(io, result->fd, io->rbuff, io->rbuff_len, 0);
        }
      } else if (result->event == EVENT_WRITE) {
        ret = entries->res;
        io = io_uring_cqe_get_data(entries);
        if (ret <= 0) {
          printf("wclose fd:%d,ret:%d\n",io->c_info.fd,ret);
          close(result->fd);
          free_io_rwfree(io);
          free(io);
        } else if (ret > 0) {
        //   printf("write\n");
          io->c_info.event = EVENT_READ;
          set_event_recv(io, result->fd, io->rbuff, io->rbuff_len, 0);
        }
      }
    }

    io_uring_cq_advance(&loop.ring, nready);
  }
}
